import os

from omegaconf import OmegaConf

from ultralytics import YOLO
from utils.mysql_crud import create_connection, update_data, select_data
from utils.redis_crud import get_redis_client, del_pause
import datetime
import uuid
from utils.oss_file import Logger, upload_local_file
from utils.set_seed import seed_torch


class TrainVal:
    def __init__(self, config, device):
        self.config = config
        self.device = device
        self.connection = create_connection()
        self.redis_client = get_redis_client()
        yolo_version = config.model.parameters.version
        yolo_config_yaml = config.model.parameters.config_yaml

        self.resume = True
        self.logger = Logger()
        if not self.resume_train():
            self.project = f'{config.save_path}/object_detect/{config.model.name}'
            self.name = f'{datetime.datetime.now().strftime("%Y-%m-%d")}-{uuid.uuid4()}'
            save_path = self.project + "/" + self.name
            self.save_path = save_path
            if not os.path.exists(save_path):
                os.makedirs(save_path)
            update_query = "UPDATE train_task SET is_pause=%s,checkpoint_folder=%s,log=%s WHERE id=%s"
            update_data(self.connection, update_query, (1, save_path, self.logger.link, self.config.train_task_id))
            self.logger.log(OmegaConf.to_yaml(config))
            self.resume = False

        self.model = YOLO(f"ultralytics/cfg/models/{yolo_version}/{yolo_config_yaml}")
        # if config.model.pretrained:
        #     self.model.load(self.save_path + f"yolo{yolo_version}{yolo_scale}.pt")

        seed_torch(config.seed)

    def train(self):
        self.model.train(data=self.config.dataset.data_folder,
                         imgsz=(self.config.dataset.height + self.config.dataset.width) // 2,  # todo
                         epochs=self.config.epoch,
                         batch=self.config.batch_size,
                         lr0=self.config.lr,
                         lrf=0.,
                         pretrained=self.config.model.pretrained,
                         device=self.device,
                         cos_lr=self.config.scheduler,
                         optimizer=self.config.optim,
                         seed=self.config.seed,
                         workers=self.config.num_workers,
                         project=self.project,
                         name=self.name,
                         exist_ok=True,
                         resume=self.resume
                         )
        self.final_crud()

    def resume_train(self):
        # 查询数据库中checkpoint_folder是否有值且status不为0
        select_query = f"SELECT checkpoint_folder,status,log FROM train_task WHERE id = {self.config.train_task_id} limit 1"
        save_path, status, log = select_data(self.connection, select_query)[0]
        if save_path is not None and save_path != '' and status > 0:
            self.save_path = save_path
            self.name = self.save_path.split('/')[-1]
            self.project = self.save_path.replace("/" + self.name, "")
            return True
        return False

    def final_crud(self):
        self.logger.log('uploading best model to oss ......')
        best_net_path = os.path.join(self.save_path, "weights", "best.pt")
        link = upload_local_file(best_net_path)

        update_query = "UPDATE train_task SET checkpoint_link=%s,status=%s,is_pause=%s WHERE id=%s"
        update_data(self.connection, update_query,
                    (link, 1, 1, self.config.train_task_id))
        del_pause(self.redis_client, self.config.train_task_id)
        self.connection.close()
        self.redis_client.close()
        self.logger.log('mysql and redis connection closed.')
